#!/usr/bin python3
# -*- coding:UTF-8 -*-
# Author: nigo


import plotly.graph_objects as go
import pandas as pd


def get_label(df):
    """
    根据两列节点及最后列数值列进行计算
    df:列数为3的dataframe表
    value:最后一列
    return:labels,parents,values
    """
    # 前两列名称列表
    groupby_name = list(df.columns[:2])
    # 按前两列聚合求和
    df = df.groupby(groupby_name,as_index=False).sum()
    # 第二列为节点列表
    labels = list(df.iloc[:,1])
    # 第一列为父节点列表
    parents = list(df.iloc[:,0])
    # 最后一列为数值列表
    values = list(df.iloc[:,-1])
    return labels,parents,values

def df2sunburst(df,top="",value=False):
    """将dataframe转换为旭日图所需数据
    df:dataframe二维表
    top:是否添加顶级节点名称，默认为空
    value:最后一列是否为数值，默认为False
    return:labels,parents,values
    """
    # 如果df为空
    if df.empty:
        return [],[],[],[]
    # 将空值用""填充
    df = df.fillna("")
    # 节点列表
    labels = []
    # 父节点列表
    parents = []
    # 值
    values = []
    # 计数器
    num = 0
    if top:
        # df在第一列插入值为top的列
        df.insert(0,'top',top)
    # else:
        # df.insert(0,'top','')
    # 若最后一列为数值
    if value:
        # 设置最后一列列名为'value'
        df.columns = list(df.columns[:-1]) + ['value'] 
    else:
        # 否则添加最后一列，列名为value,值为0
        df['value'] = 0
    # 原始数据中非数值列的列数
    num = len(df.columns) -2
    for i in range(num):
        # 循环获取相领两列以及最后数值列为sub_df
        sub_df = df.iloc[:,[i,i+1,-1]]
        # 根据这两列计算出节点列表，父节点列表，以及对应值
        tmp_labels,tmp_parents,tmp_values = get_label(sub_df)
        labels += tmp_labels
        parents += tmp_parents
        values += tmp_values
    return labels,parents,values

def df2data(df,top="",value=False):
    """将dataframe转换为桑基图所需数据
    df:dataframe二维表
    top:是否添加顶级节点名称，默认为空
    value:最后一列是否为数值，默认为False
    return:labels,sources,targets,values
    """
    # 利用df2sunbusrt计算出targets名称,sources名称,values
    targets,sources,values = df2sunburst(df,top,value)
    # 将targets和sources名称合并并去重
    labels = list(set(targets + sources))
    # 创建节点名称与索引值的对应关系字典
    label2index = {value:index for index,value in enumerate(labels)}
    # 计算出targets索引
    targets = [label2index[target] for target in targets]
    # 计算出sources索引
    sources = [label2index[source] for source in sources]
    # 返回各项值
    return labels,sources,targets,values


if __name__ == "__main__":
    df = pd.read_clipboard(delimiter='\s+',thousands=',')
    label,source,target,value = df2data(df,top="生产成本",value=True)
    
    san= go.Sankey(
        node =dict(
            label=label
        ),
        link = dict(
            source = source,
            target = target,
            value =  value,
        )
    )
    fig = go.Figure(san)
    fig.show()
